# -*- coding: utf-8 -*-
# @author: HRUN

import json
import asyncio
import logging
from datetime import datetime
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from performance.models import TaskReport, PerformanceTask

logger = logging.getLogger(__name__)


class PerformanceMonitorConsumer(AsyncWebsocketConsumer):
    """性能监控WebSocket消费者"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.room_group_name = None
        self.task_id = None
        self.monitoring = False
        
    async def connect(self):
        """WebSocket连接建立"""
        self.task_id = self.scope['url_route']['kwargs']['task_id']
        self.room_group_name = f'performance_monitor_{self.task_id}'
        
        # 加入房间组
        await self.channel_layer.group_add(
            self.room_group_name,
            self.channel_name
        )
        
        await self.accept()
        
        # 发送连接成功消息
        await self.send(text_data=json.dumps({
            'type': 'connection_established',
            'message': f'已连接到任务 {self.task_id} 的监控',
            'task_id': self.task_id,
            'timestamp': datetime.now().isoformat()
        }))
        
        logger.info(f"WebSocket连接建立 - 任务ID: {self.task_id}")

    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        if self.room_group_name:
            await self.channel_layer.group_discard(
                self.room_group_name,
                self.channel_name
            )
        
        self.monitoring = False
        logger.info(f"WebSocket连接断开 - 任务ID: {self.task_id}, 关闭码: {close_code}")

    async def receive(self, text_data):
        """接收客户端消息"""
        try:
            text_data_json = json.loads(text_data)
            message_type = text_data_json.get('type')
            
            if message_type == 'start_monitoring':
                await self.start_monitoring()
            elif message_type == 'stop_monitoring':
                await self.stop_monitoring()
            elif message_type == 'get_current_status':
                await self.send_current_status()
            else:
                await self.send(text_data=json.dumps({
                    'type': 'error',
                    'message': f'未知的消息类型: {message_type}'
                }))
                
        except json.JSONDecodeError:
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': '无效的JSON格式'
            }))
        except Exception as e:
            logger.error(f"处理WebSocket消息失败: {e}")
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'处理消息失败: {str(e)}'
            }))

    async def start_monitoring(self):
        """开始监控"""
        self.monitoring = True
        await self.send(text_data=json.dumps({
            'type': 'monitoring_started',
            'message': '开始实时监控',
            'task_id': self.task_id
        }))
        
        # 启动定期发送数据的任务
        asyncio.create_task(self.send_periodic_updates())

    async def stop_monitoring(self):
        """停止监控"""
        self.monitoring = False
        await self.send(text_data=json.dumps({
            'type': 'monitoring_stopped',
            'message': '停止实时监控',
            'task_id': self.task_id
        }))

    async def send_current_status(self):
        """发送当前状态"""
        try:
            # 获取最新的报告数据
            reports = await self.get_task_reports()
            
            if reports:
                latest_report = reports[0]
                await self.send(text_data=json.dumps({
                    'type': 'current_status',
                    'data': {
                        'task_id': self.task_id,
                        'report_id': latest_report['id'],
                        'report_name': latest_report['reportName'],
                        'status': latest_report['reportStatus'],
                        'avg_tps': latest_report['avgTps'],
                        'avg_cpu': latest_report['avgCpu'],
                        'avg_memory': latest_report['avgMemory'],
                        'total_requests': latest_report['totalRequests'],
                        'success_requests': latest_report['successRequests'],
                        'failed_requests': latest_report['failedRequests'],
                        'error_rate': latest_report['errorRate'],
                        'start_time': latest_report['startTime'].isoformat() if latest_report['startTime'] else None,
                        'duration': latest_report['duration']
                    },
                    'timestamp': datetime.now().isoformat()
                }))
            else:
                await self.send(text_data=json.dumps({
                    'type': 'current_status',
                    'data': None,
                    'message': '暂无报告数据',
                    'timestamp': datetime.now().isoformat()
                }))
                
        except Exception as e:
            logger.error(f"获取当前状态失败: {e}")
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'获取状态失败: {str(e)}'
            }))

    async def send_periodic_updates(self):
        """定期发送更新数据"""
        while self.monitoring:
            try:
                await self.send_current_status()
                await asyncio.sleep(5)  # 每5秒发送一次更新
            except Exception as e:
                logger.error(f"发送定期更新失败: {e}")
                break

    @database_sync_to_async
    def get_task_reports(self):
        """获取任务报告（异步数据库查询）"""
        try:
            reports = TaskReport.objects.filter(
                task_id=self.task_id
            ).order_by('-create_time').values(
                'id', 'reportName', 'reportStatus', 'avgTps', 'avgCpu', 
                'avgMemory', 'totalRequests', 'successRequests', 
                'failedRequests', 'errorRate', 'startTime', 'duration'
            )[:5]  # 获取最新5个报告
            return list(reports)
        except Exception as e:
            logger.error(f"获取报告数据失败: {e}")
            return []

    # 组消息处理器
    async def performance_update(self, event):
        """处理性能更新消息"""
        await self.send(text_data=json.dumps({
            'type': 'performance_update',
            'data': event['data'],
            'timestamp': event.get('timestamp', datetime.now().isoformat())
        }))

    async def test_started(self, event):
        """处理测试开始消息"""
        await self.send(text_data=json.dumps({
            'type': 'test_started',
            'message': '性能测试已开始',
            'data': event.get('data', {}),
            'timestamp': event.get('timestamp', datetime.now().isoformat())
        }))

    async def test_completed(self, event):
        """处理测试完成消息"""
        await self.send(text_data=json.dumps({
            'type': 'test_completed',
            'message': '性能测试已完成',
            'data': event.get('data', {}),
            'timestamp': event.get('timestamp', datetime.now().isoformat())
        }))

    async def test_failed(self, event):
        """处理测试失败消息"""
        await self.send(text_data=json.dumps({
            'type': 'test_failed',
            'message': '性能测试失败',
            'data': event.get('data', {}),
            'error': event.get('error', '未知错误'),
            'timestamp': event.get('timestamp', datetime.now().isoformat())
        }))


class RealtimeDataBroadcaster:
    """实时数据广播器"""
    
    @staticmethod
    async def broadcast_performance_update(task_id, data):
        """广播性能更新"""
        from channels.layers import get_channel_layer
        
        channel_layer = get_channel_layer()
        if channel_layer:
            await channel_layer.group_send(
                f'performance_monitor_{task_id}',
                {
                    'type': 'performance_update',
                    'data': data,
                    'timestamp': datetime.now().isoformat()
                }
            )

    @staticmethod
    async def broadcast_test_status(task_id, status, data=None, error=None):
        """广播测试状态"""
        from channels.layers import get_channel_layer
        
        channel_layer = get_channel_layer()
        if channel_layer:
            message = {
                'type': f'test_{status}',
                'data': data or {},
                'timestamp': datetime.now().isoformat()
            }
            
            if error:
                message['error'] = error
                
            await channel_layer.group_send(
                f'performance_monitor_{task_id}',
                message
            )


class PerformanceReportConsumer(AsyncWebsocketConsumer):
    """性能报告实时更新WebSocket消费者"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.room_group_name = None
        self.report_id = None
        self.monitoring = False
        
    async def connect(self):
        """WebSocket连接建立"""
        self.report_id = self.scope['url_route']['kwargs']['report_id']
        self.room_group_name = f'performance_report_{self.report_id}'
        
        # 加入房间组
        await self.channel_layer.group_add(
            self.room_group_name,
            self.channel_name
        )
        
        await self.accept()
        
        # 发送连接成功消息
        await self.send(text_data=json.dumps({
            'type': 'connection_established',
            'message': f'已连接到报告 {self.report_id} 的实时更新',
            'report_id': self.report_id,
            'timestamp': datetime.now().isoformat()
        }))
        
        logger.info(f"报告WebSocket连接建立 - 报告ID: {self.report_id}")

    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        if self.room_group_name:
            await self.channel_layer.group_discard(
                self.room_group_name,
                self.channel_name
            )
        
        self.monitoring = False
        logger.info(f"报告WebSocket连接断开 - 报告ID: {self.report_id}, 关闭码: {close_code}")

    async def receive(self, text_data):
        """接收客户端消息"""
        try:
            text_data_json = json.loads(text_data)
            message_type = text_data_json.get('type')
            
            if message_type == 'start_monitoring':
                await self.start_monitoring()
            elif message_type == 'stop_monitoring':
                await self.stop_monitoring()
            elif message_type == 'get_report_data':
                await self.send_report_data()
            elif message_type == 'get_system_status':
                await self.send_system_status()
            else:
                await self.send(text_data=json.dumps({
                    'type': 'error',
                    'message': f'未知的消息类型: {message_type}'
                }))
                
        except json.JSONDecodeError:
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': '无效的JSON格式'
            }))
        except Exception as e:
            logger.error(f"处理WebSocket消息失败: {e}")
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'处理消息失败: {str(e)}'
            }))

    async def start_monitoring(self):
        """开始监控"""
        self.monitoring = True
        await self.send(text_data=json.dumps({
            'type': 'monitoring_started',
            'message': '开始实时监控报告',
            'report_id': self.report_id
        }))
        
        # 启动定期发送数据的任务
        asyncio.create_task(self.send_periodic_updates())

    async def stop_monitoring(self):
        """停止监控"""
        self.monitoring = False
        await self.send(text_data=json.dumps({
            'type': 'monitoring_stopped',
            'message': '停止实时监控报告',
            'report_id': self.report_id
        }))

    async def send_report_data(self):
        """发送报告数据"""
        try:
            report_data = await self.get_report_data()
            
            if report_data:
                await self.send(text_data=json.dumps({
                    'type': 'report_data',
                    'data': report_data,
                    'timestamp': datetime.now().isoformat()
                }))
            else:
                await self.send(text_data=json.dumps({
                    'type': 'report_data',
                    'data': None,
                    'message': '报告不存在',
                    'timestamp': datetime.now().isoformat()
                }))
                
        except Exception as e:
            logger.error(f"获取报告数据失败: {e}")
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'获取报告数据失败: {str(e)}'
            }))

    async def send_system_status(self):
        """发送系统状态"""
        try:
            from performanceengine.systemMonitor import get_current_system_monitor
            
            monitor = get_current_system_monitor()
            if monitor:
                system_info = monitor.collect_system_info()
                await self.send(text_data=json.dumps({
                    'type': 'monitoring_update',
                    'data': system_info,
                    'timestamp': datetime.now().isoformat()
                }))
            else:
                # 获取基本系统信息
                system_info = await self.get_basic_system_info()
                await self.send(text_data=json.dumps({
                    'type': 'monitoring_update',
                    'data': system_info,
                    'timestamp': datetime.now().isoformat()
                }))
                
        except Exception as e:
            logger.error(f"获取系统状态失败: {e}")
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'获取系统状态失败: {str(e)}'
            }))

    async def send_logs_update(self):
        """发送日志更新"""
        try:
            # 生成模拟日志数据
            logs = await self.generate_sample_logs()
            
            for log in logs:
                await self.send(text_data=json.dumps({
                    'type': 'log_update',
                    'data': log,
                    'timestamp': datetime.now().isoformat()
                }))
                
        except Exception as e:
            logger.error(f"发送日志更新失败: {e}")

    async def send_periodic_updates(self):
        """定期发送更新数据"""
        while self.monitoring:
            try:
                await self.send_report_data()
                await self.send_system_status()
                await asyncio.sleep(10)  # 每10秒发送一次更新
            except Exception as e:
                logger.error(f"发送定期更新失败: {e}")
                break

    @database_sync_to_async
    def get_report_data(self):
        """获取报告数据（异步数据库查询）"""
        try:
            from performance.models import TaskReport
            
            report = TaskReport.objects.select_related('task', 'env').get(id=self.report_id)
            
            return {
                'id': report.id,
                'reportName': report.reportName,
                'reportStatus': report.reportStatus,
                'avgTps': report.avgTps,
                'maxTps': report.maxTps,
                'minTps': report.minTps,
                'avgCpu': report.avgCpu,
                'avgMemory': report.avgMemory,
                'avgResponseTime': report.avgResponseTime,
                'maxResponseTime': report.maxResponseTime,
                'minResponseTime': report.minResponseTime,
                'p50ResponseTime': report.p50ResponseTime,
                'p90ResponseTime': report.p90ResponseTime,
                'p95ResponseTime': report.p95ResponseTime,
                'p99ResponseTime': report.p99ResponseTime,
                'totalRequests': report.totalRequests,
                'successRequests': report.successRequests,
                'failedRequests': report.failedRequests,
                'errorRate': report.errorRate,
                'maxUsers': report.maxUsers,
                'avgUsers': report.avgUsers,
                'startTime': report.startTime.isoformat() if report.startTime else None,
                'endTime': report.endTime.isoformat() if report.endTime else None,
                'duration': report.duration,
                'executor': report.executor,
                'create_time': report.create_time.isoformat() if report.create_time else None,
                'task': {
                    'id': report.task.id,
                    'taskName': report.task.taskName,
                    'taskType': report.task.taskType,
                    'runPattern': report.task.runPattern,
                    'distributed_mode': report.task.distributed_mode
                } if report.task else None,
                'env': {
                    'id': report.env.id,
                    'name': report.env.name,
                    'host': report.env.host
                } if report.env else None
            }
        except TaskReport.DoesNotExist:
            return None
        except Exception as e:
            logger.error(f"获取报告数据失败: {e}")
            return None

    async def get_basic_system_info(self):
        """获取基本系统信息"""
        try:
            import psutil
            import random
            
            return {
                'cpu_percent': psutil.cpu_percent(interval=1),
                'memory_percent': psutil.virtual_memory().percent,
                'disk_percent': psutil.disk_usage('/').percent,
                'network_sent': psutil.net_io_counters().bytes_sent,
                'network_recv': psutil.net_io_counters().bytes_recv,
                'active_connections': len(psutil.net_connections()),
                'current_rps': random.randint(50, 200),
                'current_users': random.randint(10, 100),
                'error_rate': random.uniform(0, 5),
                'server_type': 'single',
                'cpu_cores': psutil.cpu_count(),
                'total_memory': psutil.virtual_memory().total,
                'uptime': None
            }
        except ImportError:
            # 如果psutil不可用，返回模拟数据
            import random
            return {
                'cpu_percent': random.randint(30, 80),
                'memory_percent': random.randint(40, 85),
                'disk_percent': random.randint(20, 70),
                'network_sent': random.randint(1000000, 10000000),
                'network_recv': random.randint(1000000, 10000000),
                'active_connections': random.randint(50, 200),
                'current_rps': random.randint(50, 200),
                'current_users': random.randint(10, 100),
                'error_rate': random.uniform(0, 5),
                'server_type': 'single',
                'cpu_cores': 4,
                'total_memory': 8589934592,  # 8GB
                'uptime': None
            }

    async def generate_sample_logs(self):
        """生成示例日志"""
        import random
        from datetime import timedelta
        
        levels = ['info', 'warning', 'error']
        messages = [
            '用户连接建立成功',
            '发送HTTP请求',
            '接收到响应数据',
            '更新性能统计',
            '检测到高响应时间',
            '错误率轻微上升'
        ]
        
        logs = []
        now = datetime.now()
        
        # 随机生成1-3条日志
        for i in range(random.randint(1, 3)):
            logs.append({
                'timestamp': (now - timedelta(seconds=random.randint(0, 30))).isoformat(),
                'level': random.choice(levels),
                'message': random.choice(messages),
                'report_id': self.report_id
            })
        
        return logs

    # 组消息处理器
    async def performance_update(self, event):
        """处理性能更新消息"""
        await self.send(text_data=json.dumps({
            'type': 'performance_update',
            'data': event['data'],
            'timestamp': event.get('timestamp', datetime.now().isoformat())
        }))

    async def log_update(self, event):
        """处理日志更新消息"""
        await self.send(text_data=json.dumps({
            'type': 'log_update',
            'data': event['data'],
            'timestamp': event.get('timestamp', datetime.now().isoformat())
        }))

    async def monitoring_update(self, event):
        """处理监控更新消息"""
        await self.send(text_data=json.dumps({
            'type': 'monitoring_update',
            'data': event['data'],
            'timestamp': event.get('timestamp', datetime.now().isoformat())
        }))


class ReportDataBroadcaster:
    """报告数据广播器"""
    
    @staticmethod
    async def broadcast_performance_update(report_id, data):
        """广播性能更新到报告级别"""
        from channels.layers import get_channel_layer
        
        channel_layer = get_channel_layer()
        if channel_layer:
            await channel_layer.group_send(
                f'performance_report_{report_id}',
                {
                    'type': 'performance_update',
                    'data': data,
                    'timestamp': datetime.now().isoformat()
                }
            )

    @staticmethod
    async def broadcast_log_update(report_id, log_data):
        """广播日志更新"""
        from channels.layers import get_channel_layer
        
        channel_layer = get_channel_layer()
        if channel_layer:
            await channel_layer.group_send(
                f'performance_report_{report_id}',
                {
                    'type': 'log_update',
                    'data': log_data,
                    'timestamp': datetime.now().isoformat()
                }
            )

    @staticmethod
    async def broadcast_monitoring_update(report_id, monitoring_data):
        """广播监控更新"""
        from channels.layers import get_channel_layer
        
        channel_layer = get_channel_layer()
        if channel_layer:
            await channel_layer.group_send(
                f'performance_report_{report_id}',
                {
                    'type': 'monitoring_update',
                    'data': monitoring_data,
                    'timestamp': datetime.now().isoformat()
                }
            )